#!/usr/bin/env python
# -*- coding:utf-8 -*-
# 将一个文件夹下图片按比例分在三个文件夹下
import os
import random
import shutil
from shutil import copy2
from tqdm import tqdm

datadir_normal = "G:\数据集\植保所\分类目标图片\other"
save_path="G:\数据集\植保所\分类目标图片\斑潜蝇\深度学习数据集\other"

all_data_img = os.listdir(datadir_normal)  # （图片文件夹）

num_all_data = len(all_data_img)
print("num_all_data: " + str(num_all_data))
index_list = list(range(num_all_data))
print(index_list)
random.shuffle(index_list)
num = 0

trainDir = os.path.join(save_path,"train")  # （将训练集放在这个文件夹下）
if not os.path.exists(trainDir):
    os.mkdir(trainDir)

validDir = os.path.join(save_path,"val")  # （将验证集放在这个文件夹下）
if not os.path.exists(validDir):
    os.mkdir(validDir)

testDir = os.path.join(save_path,"test")  # （将测试集放在这个文件夹下）
if not os.path.exists(testDir):
    os.mkdir(testDir)

for i in tqdm(index_list):
    fileName = os.path.join(datadir_normal, all_data_img[i])

    if num < num_all_data * 0.7:
        copy2(fileName, trainDir)

    elif num > num_all_data * 0.7 and num < num_all_data * 0.8:
        # print(str(fileName))
        copy2(fileName, validDir)

    else:
        copy2(fileName, testDir)

    num += 1
